package qy.jalgotrade.utils.plotter;

import java.lang.reflect.InvocationTargetException;
import java.time.ZonedDateTime;
import java.time.format.DateTimeFormatter;
import java.util.*;

import jep.*;
import org.apache.commons.collections4.MapUtils;
import org.apache.commons.collections4.set.ListOrderedSet;
import org.apache.commons.io.IOUtils;
import org.apache.commons.lang3.exception.ExceptionUtils;
import org.apache.commons.lang3.tuple.ImmutableTriple;
import org.apache.commons.lang3.tuple.Triple;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import jep.python.PyCallable;
import jep.python.PyObject;
import qy.jalgotrade.Initializer;
import qy.jalgotrade.bar.Bar;
import qy.jalgotrade.bar.Bars;
import qy.jalgotrade.broker.Broker;
import qy.jalgotrade.broker.OrderEvent;
import qy.jalgotrade.event.EventHandler;
import qy.jalgotrade.strategy.BaseStrategy;
import qy.jalgotrade.utils.CommonUtils;

/**
 * Class responsible for plotting a strategy execution.
 * 
 * @author c-geo
 *
 */
public class StrategyPlotter implements EventHandler {

	private static final Logger logger = LoggerFactory.getLogger(StrategyPlotter.class);

	private static final String PLOTTER_PY = "py/plotter.py";

	private static final String PLOTTER_PY_ENCODING = "UTF-8";

	private static Jep jepInterp;

	static {

		try {
			// init jep:
			JepConfig config = new JepConfig();
			config.setRedirectOutputStreams(true);
			SharedInterpreter.setConfig(config);
			jepInterp = new SharedInterpreter();
			String plotterPy = IOUtils.toString(
					Objects.requireNonNull(StrategyPlotter.class.getClassLoader().getResourceAsStream(PLOTTER_PY)),
			        PLOTTER_PY_ENCODING);
			jepInterp.exec(plotterPy);
			// init jep log:
			String jepLogFile = String.format("%s%s%s", System.getProperty(Initializer.LOG_DIR),
					System.getProperty("file.separator"),
					Initializer.APP_DEPLOY_CONFIG.getSection("jep").getString("logFilename"));
			String logLevel = Initializer.APP_DEPLOY_CONFIG.getSection("jep").getString("logLevel");
			String logFormat = Initializer.APP_DEPLOY_CONFIG.getSection("jep").getString("logFormat");
			jepInterp.exec(String.format("jepLogFile = \"%s\"; logLevel = \"%s\"; logFormat = \"%s\"; initLogging()",
					jepLogFile, logLevel, logFormat));
		} catch (Exception e) {
			logger.error("Exception occured initializing Jep Interpreter: {}", ExceptionUtils.getStackTrace(e));
		}
	}

	/**
	 * strategy 的主 frequency:
	 */
	private final Bar.Frequency __frequency;

	/**
	 * 确保所有 subplots 横轴 (x 轴) 统一为 datetime:
	 */
	private final Set<ZonedDateTime> __dateTimes;

	/**
	 * 是否默认生成所有 instruments 的 subplot:
	 */
	private final boolean __plotAllInstruments;

	/**
	 * 是否标记 buy/sell (broker 产生的 成交事件):
	 */
	private final boolean __plotBuySell;

	/**
	 * instrument 所对应的 subplots (qyalgotrader.utils.myplotter.InstrumentSubplot):
	 */
	private final Map<String, InstrumentSubplot> __barSubplots;

	/**
	 * 自定义 name 的 subplots:
	 */
	private final Map<String, Subplot> __namedSubplots;

	/**
	 * portfolio 总市值 subplot:
	 */
	private Subplot __portfolioSubplot;

	/**
	 * MyCursor 引用, 防止其被 GC:
	 */
	private PyObject __myCursor;

	/**
	 *
	 */
	private boolean __figImplReady;

	/**
	 * 
	 * @return Interpreter
	 */
	public static Interpreter jepInterpreter() {

		return jepInterp;
	}

	/**
	 * 
	 * @param strat The strategy to plot.
	 */
	public StrategyPlotter(BaseStrategy strat) {

		this(strat, true, true, true);
	}

	/**
	 *
	 * @param strat The strategy to plot.
	 * @param plotAllInstrumentsP default: true
	 * @param plotBuySellP default: true
	 * @param plotPortfolioP default: true
	 */
	public StrategyPlotter(BaseStrategy strat, Optional<Boolean> plotAllInstrumentsP, Optional<Boolean> plotBuySellP,
						   Optional<Boolean> plotPortfolioP) {

		this(strat, plotAllInstrumentsP.orElse(true), plotBuySellP.orElse(true),
				plotPortfolioP.orElse(true));
	}

	/**
	 * 
	 * @param strat              The strategy to plot.
	 * @param plotAllInstruments Set to True to get a subplot for each instrument available.
	 * @param plotBuySell        Set to True to get the buy/sell events plotted for each instrument
	 *                           available.
	 * @param plotPortfolio      Set to True to get the portfolio value (shares + cash) plotted.
	 */
	public StrategyPlotter(BaseStrategy strat, boolean plotAllInstruments, boolean plotBuySell, boolean plotPortfolio) {

		__frequency = strat.getFeed().getFrequency();
		__dateTimes = new ListOrderedSet<>();
		__plotAllInstruments = plotAllInstruments;
		__plotBuySell = plotBuySell;
		__barSubplots = MapUtils.orderedMap(new HashMap<>());
		__namedSubplots = MapUtils.orderedMap(new HashMap<>());
		__portfolioSubplot = null;
		if (plotPortfolio) {
			__portfolioSubplot = new Subplot(__frequency);
			__portfolioSubplot.setName("Portfolio");
		}
		__myCursor = null;
		__figImplReady = false;

		// 订阅 strategy 的 bar-processed 事件 需早于 broker 的 order-updated 事件 (成交 晚于 信号):
		strat.getBarsProcessedEvent().subscribe(this);
		strat.getBroker().getOrderUpdatedEvent().subscribe(this);
	}

	/* (non-Javadoc)
	 * @see qy.jalgotrade.event.EventHandler#handle(java.lang.String, java.util.Map)
	 */
	@Override
	public void handle(String eventName, Map<String, Object> params) throws Exception {

		// TODO: Event 分类
		if ("strat-bar-proced-event".equals(eventName)) {
			__onBarsProcessed((BaseStrategy) params.get("strategy"), (Bars) params.get("bars"));
		} else if ("broker-order-event".equals(eventName)) {
			__onOrderEvent((Broker) params.get("broker"), (OrderEvent) params.get("orderEvent"));
		} else {
			throw new IllegalArgumentException("Unknown eventName: " + eventName);
		}
	}

	/**
	 * Returns the InstrumentSubplot for a given instrument. 创建 (若不存在) 并返回指定 instrument 的 subplot:
	 * 
	 * @param instrument instrument
	 * @return InstrumentSubplot
	 */
	public InstrumentSubplot getInstrumentSubplot(String instrument) {

		InstrumentSubplot ret = __barSubplots.get(instrument);
		if (ret == null) {
			ret = new InstrumentSubplot(instrument, __frequency, __plotBuySell);
			ret.setName(instrument);
			__barSubplots.put(instrument, ret);
		}
		return ret;
	}

	/**
	 * Returns a Subplot by name. If the subplot doesn't exist, it gets created. 创建 (若不存在) 并返回 自定义 name
	 * 的 subplot:
	 * 
	 * @param name The name of the Subplot to get or create.
	 * @return Subplot
	 */
	public Subplot getOrCreateSubplot(String name) {

		Subplot ret = __namedSubplots.get(name);
		if (ret == null) {
			ret = new Subplot(__frequency);
			ret.setName(name);
			__namedSubplots.put(name, ret);
		}
		return ret;
	}

	/**
	 * Returns the subplot where the portfolio values get plotted.
	 * 
	 * @return Subplot
	 */
	public Subplot getPortfolioSubplot() {

		return __portfolioSubplot;
	}

	private void __checkCreateInstrumentSubplot(String instrument) {

		if (!__barSubplots.containsKey(instrument)) {
			getInstrumentSubplot(instrument);
		}
	}

	private void __onBarsProcessed(BaseStrategy strat, Bars bars) {

		ZonedDateTime dateTime = bars.getDateTime();
		__dateTimes.add(dateTime);

		if (__plotAllInstruments) {
			for (String instrument : bars.getInstruments()) {
				__checkCreateInstrumentSubplot(instrument);
			}
		}

		// Notify named subplots.
		for (Subplot subplot : __namedSubplots.values()) {
			subplot.onBars(bars);
		}

		// Notify bar subplots.
		for (InstrumentSubplot subplot : __barSubplots.values()) {
			subplot.onBars(bars);
		}

		// Feed the portfolio evolution subplot.
		if (__portfolioSubplot != null) {
			try {
				__portfolioSubplot.getSeries("Portfolio").addValue(dateTime, strat.getBroker().getEquity());
				// This is in case additional dataseries were added to the portfolio subplot.
				__portfolioSubplot.onBars(bars);
			} catch (NoSuchMethodException | IllegalAccessException | InvocationTargetException
			        | InstantiationException e) {

				logger.error("Exception occured creating Series: {}", ExceptionUtils.getStackTrace(e));
				throw new IllegalStateException(e);
			}
		}
	}

	private void __onOrderEvent(Broker broker, OrderEvent orderEvent) {

		// Notify BarSubplots
		for (InstrumentSubplot subplot : __barSubplots.values()) {
			subplot.onOrderEvent(broker, orderEvent);
		}
	}

	private Triple<PyObject, List<PyObject>, PyObject> __buildFigureImpl(ZonedDateTime fromDateTime,
	        ZonedDateTime toDateTime) throws JepException {

		List<ZonedDateTime> dateTimes = new ArrayList<>();
		for (ZonedDateTime e : __dateTimes) {
			if (toDateTime != null && e.compareTo(toDateTime) > 0
			        || fromDateTime != null && e.compareTo(fromDateTime) < 0) {
				continue;
			}
			dateTimes.add(e);
		}
		Collections.sort(dateTimes);

		List<Subplot> subplots = new ArrayList<>();
		List<String> subplotNames = new ArrayList<>();

		// {subplotName: [subplot, mplSubplot]}:
		Map<String, List<Object>> subplotsMap = new HashMap<>();
		for (Map.Entry<String, InstrumentSubplot> e : __barSubplots.entrySet()) {
			subplotsMap.put(e.getKey(), Arrays.asList(e.getValue(), null));
			subplotNames.add(e.getKey());
			subplots.add(e.getValue());
		}
		for (Map.Entry<String, Subplot> e : __namedSubplots.entrySet()) {
			subplotsMap.put(e.getKey(), Arrays.asList(e.getValue(), null));
			subplotNames.add(e.getKey());
			subplots.add(e.getValue());
		}

		if (__portfolioSubplot != null) {
			subplotsMap.put("Portfolio", Arrays.asList(__portfolioSubplot, null));
			subplotNames.add("Portfolio");
			subplots.add(__portfolioSubplot);
		}

		// add shared object to jep interpreter:
		PyObject plotterDateTimes = CommonUtils.createPyList(jepInterp, "plotterDateTimes");
		PyObject plotterDateTimeIdx = CommonUtils.createPyList(jepInterp, "plotterDateTimeIdx");
		for (int i = 0; i < dateTimes.size(); i++) {
			CommonUtils.pyInvoke(plotterDateTimes, "append",
			        dateTimes.get(i).format(DateTimeFormatter.ofPattern("yyyyMMdd HH:mm")));
			CommonUtils.pyInvoke(plotterDateTimeIdx, "append", i);
		}

		// Build each subplot (qyalgotrader.utils.myplotter.Subplot).
		// fig, axes = plt.subplots(nrows=len(subplots), sharex=True, squeeze=False)
		String createSubplotsPy = String.format(
		        "fig, axes = plt.subplots(nrows=%s, sharex=True, squeeze=False); axesLs = axes.flatten().tolist()",
		        subplots.size());
		jepInterp.exec(createSubplotsPy);

		PyObject fig = jepInterp.getValue("fig", PyObject.class);
		PyObject axesLs = jepInterp.getValue("axesLs", PyObject.class);
		PyObject mplSubplots = CommonUtils.createPyList(jepInterp,
		        CommonUtils.genPyObjectName(getClass().getName(), "__buildFigureImpl", "mplSubplots"));
		List<PyObject> rtMplSubplots = new ArrayList<>();
		for (int i = 0; i < subplots.size(); i++) {

			if (!subplots.get(i).isEmpty()) {

				PyObject axesSubplot = CommonUtils.pyInvokeAs(PyObject.class, axesLs, "__getitem__", i);
				CommonUtils.pyInvoke(mplSubplots, "append", axesSubplot);
				rtMplSubplots.add(axesSubplot);
				/*
				 * Subplot#plot()
				 *     -> Series#plot()
				 *         -> mplSubplot#plot(dateTimes, values, ...)
				 */
				subplots.get(i).plot(axesSubplot, dateTimes); // qyalgotrader.utils.myplotter.Subplot#plot()
				// 开启网格线:
				// axesSubplot.grid(b=True, linestyle="--")
				// @formatter:off
				Map<String, Object> kwargs = MapUtils.putAll(new HashMap<String, Object>(), new Object[][] {
						{ "b", true },
						{ "linestyle", "--" }
					});
				// @formatter:on
				CommonUtils.pyInvoke(axesSubplot, "grid", kwargs);
			}
		}

		// 为 subplot 添加 十字光标:
		// mplCursor = MyMultiCursor(fig.canvas, mplSubplots, useblit=True, horizOn=False, vertOn=True, color="gray", linestyle="-.", linewidth=1.0)
		Object[] args = new Object[] { fig.getAttr("canvas", PyObject.class), mplSubplots };
		// @formatter:off
		Map<String, Object> kwargs = MapUtils.putAll(new HashMap<String, Object>(), new Object[][] {
				{ "useblit", true },
				{ "horizOn", false },
				{ "vertOn", true },
				{ "color", "gray" },
				{ "linestyle", "-." },
				{ "linewidth", 1.0 },
			});
		// @formatter:on
		PyObject mplCursor = jepInterp.getValue("MyMultiCursor", PyCallable.class).callAs(PyObject.class, args, kwargs);
		// 保存 MyMultiCursor 的引用于 Jep 实例中, 避免被 GC:
		jepInterp.set("myMplCursor", mplCursor);
		__myCursor = mplCursor;
		__figImplReady = true;

		return new ImmutableTriple<>(fig, rtMplSubplots, mplCursor);
	}

	/**
	 * Build a matplotlib.figure.Figure with the subplots. Must be called after running the strategy.
	 * 
	 * @return A 2 element tuple with matplotlib.figure.Figure and subplots.
	 * @throws JepException JepException
	 */
	public Triple<PyObject, List<PyObject>, PyObject> buildFigureAndSubplots() throws JepException {

		return buildFigureAndSubplots(null, null);
	}

	/**
	 * Build a matplotlib.figure.Figure with the subplots. Must be called after running the strategy.
	 * 
	 * @param fromDateTime fromDateTime An optional starting datetime.datetime. Everything before it
	 *                     won't get plotted.
	 * @return A 2 element tuple with matplotlib.figure.Figure and subplots.
	 * @throws JepException JepException
	 */
	public Triple<PyObject, List<PyObject>, PyObject> buildFigureAndSubplots(ZonedDateTime fromDateTime)
	        throws JepException {

		return buildFigureAndSubplots(fromDateTime, null);
	}

	/**
	 * Build a matplotlib.figure.Figure with the subplots. Must be called after running the strategy.
	 * 
	 * @param fromDateTime An optional starting datetime.datetime. Everything before it won't get
	 *                     plotted.
	 * @param toDateTime   An optional ending datetime.datetime. Everything after it won't get plotted.
	 * @return A 2 element tuple with matplotlib.figure.Figure and subplots.
	 * @throws JepException JepException
	 */
	public Triple<PyObject, List<PyObject>, PyObject> buildFigureAndSubplots(ZonedDateTime fromDateTime,
	        ZonedDateTime toDateTime) throws JepException {

		Triple<PyObject, List<PyObject>, PyObject> ret = __buildFigureImpl(fromDateTime, toDateTime);
		// fig.autofmt_xdate()
		CommonUtils.pyInvoke(ret.getLeft(), "autofmt_xdate");
		return ret;
	}

	/**
	 * Plot the strategy execution. Must be called after running the strategy.
	 * 
	 * @throws JepException JepException
	 */
	public void plot() throws JepException {

		plot(null, null);
	}

	/**
	 * lot the strategy execution. Must be called after running the strategy.
	 * 
	 * @param fromDateTime An optional starting datetime.datetime. Everything before it won't get
	 *                     plotted.
	 * @throws JepException JepException
	 */
	public void plot(ZonedDateTime fromDateTime) throws JepException {

		plot(fromDateTime, null);
	}

	/**
	 * Plot the strategy execution. Must be called after running the strategy.
	 * 
	 * @param fromDateTime An optional starting datetime.datetime. Everything before it won't get
	 *                     plotted.
	 * @param toDateTime   An optional ending datetime.datetime. Everything after it won't get plotted.
	 * @throws JepException JepException
	 */
	public void plot(ZonedDateTime fromDateTime, ZonedDateTime toDateTime) throws JepException {

		Triple<PyObject, List<PyObject>, PyObject> ret = __buildFigureImpl(fromDateTime, toDateTime);
		// fig.autofmt_xdate()
		CommonUtils.pyInvoke(ret.getLeft(), "autofmt_xdate");
		// plt.show()
		jepInterp.exec("plt.show()");
	}

	/**
	 * Plot the strategy execution into a file. Must be called after running the strategy.
	 *
	 * @param filename filename
	 * @param dpiP dpiP default: null
	 * @param fmtP fmtP default: null
	 * @param fromDateTimeP fromDateTimeP default: null
	 * @param toDateTimeP toDateTimeP default: null
	 * @throws JepException JepException
	 */
	public void savePlot(String filename, Optional<String> dpiP, Optional<String> fmtP,
	                     Optional<ZonedDateTime> fromDateTimeP, Optional<ZonedDateTime> toDateTimeP)
			throws JepException {

		String dpi = dpiP.orElse(null);
		String fmt = fmtP.orElse("png");
		ZonedDateTime fromDateTime = fromDateTimeP.orElse(null);
		ZonedDateTime toDateTime = toDateTimeP.orElse(null);
		savePlot(filename, dpi, fmt, fromDateTime, toDateTime);
	}

	/**
	 * Plot the strategy execution into a file. Must be called after running the strategy.
	 * 
	 * @param filename     The filename.
	 * @param dpi          The resolution in dots per inch.
	 * @param fmt          The file extension.
	 * @param fromDateTime An optional starting datetime.datetime. Everything before it won't get
	 *                     plotted.
	 * @param toDateTime   An optional ending datetime.datetime. Everything after it won't get plotted.
	 * @throws JepException JepException
	 */
	public void savePlot(String filename, String dpi, String fmt, ZonedDateTime fromDateTime, ZonedDateTime toDateTime)
	        throws JepException {

		Triple<PyObject, List<PyObject>, PyObject> ret = __buildFigureImpl(fromDateTime, toDateTime);
		// fig.autofmt_xdate()
		CommonUtils.pyInvoke(ret.getLeft(), "autofmt_xdate");
		// fig.savefig(filename, dpi=dpi, bbox_inches="tight", format=fmt)
		Object[] args = new Object[] {filename};
		// @formatter:off
		Map<String, Object> kwargs = MapUtils.putAll(new HashMap<String, Object>(), new Object[][] {
				{"dpi", dpi},
				{"bbox_inches", "tight"},
				{"format", fmt}
		});
		// @formatter:on
		CommonUtils.pyInvoke(ret.getLeft(), "savefig", args, kwargs);
	}
}
